from .loss import so3_log
from .model import ur5_fk_position_rotation
import casadi as ca
import numpy as np
import socket
import pickle

class MPCSolver:
    def __init__(self, horizon):
        self.horizon = horizon
        self.nq = 6
        self._init_mpc_solver()
        self.est_all_q = None
        pass

    def _init_mpc_solver(self):
        nq = 6
        q_vars = ca.SX.sym('q', nq * self.horizon)               # 控制变量
        q_current = ca.SX.sym('q_current', nq)              # 当前关节角
        R_target = ca.SX.sym('R_target', (3,3))            # 目标姿态（3x3 拉平）
        p_target = ca.SX.sym('p_target', 3)                 # 目标位置

        # cost function 
        q_seq = [q_vars[i*nq:(i+1)*nq] for i in range(self.horizon)]
        cost = 0
        for i in range(self.horizon):
            pos_i, R_i = ur5_fk_position_rotation(q_seq[i])
            pos_err = pos_i - p_target
            R_err = ca.mtimes([R_target, R_i])
            log_R = so3_log(R_err)  
            cost += 0.2 * ca.dot(log_R, log_R)
            cost += ca.sumsqr(pos_err)

        # 相对当前平滑
        cost += 10*ca.sumsqr(q_current - q_seq[0])
        # 平滑项
        for i in range(1, self.horizon):
            dq = q_seq[i] - q_seq[i-1]
            cost +=  0.1 * ca.sumsqr(dq)

        cost += 0.1 * ca.sumsqr(q_seq[0] - q_current)
        
        param = ca.vertcat(q_current, ca.reshape(R_target,9,1), p_target)
        nlp = {'x': q_vars, 'f': cost, 'p': param}

        self.solver = ca.nlpsol('solver', 'ipopt', nlp, {
            'ipopt.max_iter': 30,
            'ipopt.tol': 1e-2,
            'ipopt.print_level': 0,      # 0 表示无输出
            'print_time': False,         # 不打印 solver 时间
            'verbose': False             # CasADi 自身也安静
        })
    
    def solve(self, current_q, target_pos, target_R):
        p_val = np.concatenate([
            current_q,
            target_R.reshape(9),
            target_pos
        ])
        x0 = np.tile(current_q, self.horizon)
        lbx = [-2*np.pi] * self.nq * self.horizon
        ubx = [2*np.pi] * self.nq * self.horizon
        sol = self.solver(x0=x0, p=p_val, lbx=lbx, ubx=ubx)
        q_est = sol['x'].full().flatten()
        self.est_all_q = q_est
        return q_est[:self.nq]


# if __name__ == "__main__":
#     fk_func = make_fk_function()
#     initial_q = np.array([-1.611, -1.764, 1.798, -1.657, -1.577, -3.179])
#     # initial_q = np.array([-1.335, -1.624, 1.922, -1.924, -1.581, -2.903])
#     initial_pos, initial_R = fk_func(initial_q)
#     print(initial_pos)
#     print(initial_R)
#     # print([ 0.25150052, -0.41346622,  1.1327074 ])
#     q = np.array(initial_q)
#     # exit()
#     # target_pos = np.array([ -0.25150052, 0.41346622,  1.1327074 ])
#     target_pos = np.array([-0.251356, 0.413496, 0.6])
#     target_R = np.array(
#                     [[-1.000, 0.001, 0.004],
#                     [0.002, 1.000, 0.026],
#                     [-0.004, 0.026, -1.000]])
#     for t in range(50):
#         # 调用控制器
#         q_next = mpc_controller_pose(q, target_pos, target_R)

#         # 发送 q_next 到机器人（或模拟器）
#         q = q_next  # 这行在仿真中应是接收执行反馈

#         # 可选：检查末端误差
#         pos, R = fk_func(q)
#         print(f"Step {t}: pos error = {np.linalg.norm(pos.full().flatten() - np.array(target_pos)):.4f}")

# if __name__ == "__main__":
#     ctx = zmq.Context()
#     socket = ctx.socket(zmq.REP)
#     socket.bind("tcp://127.0.0.1:5555")

#     print("[MPC Server] Listening on port 5555")

#     while True:
#         msg = socket.recv()
#         data = pickle.loads(msg)

#         current_q = np.array(data["current_q"])
#         target_pos = np.array(data["target_pos"])
#         target_R = np.array(data["target_R"]).reshape(3, 3)

#         q_next = mpc_controller_pose(current_q, target_pos, target_R)
#         socket.send(pickle.dumps({"q_next": q_next.tolist()}))